package cn._2dland.websocket.model.handshake;

import java.io.IOException;
import java.io.OutputStream;
import java.security.MessageDigest;
import java.util.LinkedHashMap;
import java.util.Map;

import cn._2dland.util.Base64;

/**
 * WebSocket握手响应
 * 
 * @author deadblue
 */
public class HandShakeResponse {

	/** 计算相应值时必须使用的UUID */
	private static final String SERVER_UUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
	/** 行结束符 */
	private static final String CRLF = "\r\n";
	/** 响应状态 */
	private static final String RESPONSE_STATE = "HTTP/1.1 101 WebSocket Protocol Handshake";

	/** 响应头：升级？ */
	public static final String HEADER_UPGRADE = "Upgrade";
	/** 响应头：连接？ */
	public static final String HEADER_CONNECTION = "Connection";
	/** 响应头：响应密文 */
	public static final String HEADER_WEBSOCKET_ACCEPT = "Sec-WebSocket-Accept";
	/** 响应头：响应协议 */
	public static final String HEADER_WEBSOCKET_PROTOCOL = "Sec-WebSocket-Protocol";
	/** 响应头：来源认证 */
	public static final String HEADER_WEBSOCKET_ORIGIN = "Sec-WebSocket-Origin";
	/** 响应头：WS地址 */
	public static final String HEADER_WEBSOCKET_LOCATION = "Sec-WebSocket-Location";

	/** 响应头 */
	private Map<String, String> headers = new LinkedHashMap<String, String>(); 
	/** 负载内容 */
	private byte[] payload = null;

	/**
	 * 构造握手响应
	 * 
	 * @param request 握手请求对象
	 */
	public HandShakeResponse(HandShakeRequest request) {
		if(request.isDraft()) {
			this.makeResponseForDraft(request);
		} else {
			this.makeResponseForVer13(request);
		}
	}

	/**
	 * 为草稿版握手请求创建应答
	 * 
	 * @param request
	 */
	private void makeResponseForDraft(HandShakeRequest request) {
		// 标准头部
		headers.put(HEADER_UPGRADE, "WebSocket");
		headers.put(HEADER_CONNECTION, "Upgrade");
		headers.put(HEADER_WEBSOCKET_ORIGIN, request.getOrigin());
		headers.put(HEADER_WEBSOCKET_LOCATION, request.getLocation());

		try {
			byte[] buf = new byte[16];

			long key1 = calcKey(request.getHeader(HandShakeRequest.HEADER_KEY1));
			buf[0] = (byte)(0xff & (key1 >> 24));
			buf[1] = (byte)(0xff & (key1 >> 16));
			buf[2] = (byte)(0xff & (key1 >> 8));
			buf[3] = (byte)(0xff & key1);
			long key2 = calcKey(request.getHeader(HandShakeRequest.HEADER_KEY2));
			buf[4] = (byte)(0xff & (key2 >> 24));
			buf[5] = (byte)(0xff & (key2 >> 16));
			buf[6] = (byte)(0xff & (key2 >> 8));
			buf[7] = (byte)(0xff & key2);
			System.arraycopy(request.getPayload(), 0, buf, 8, 8);

			// 计算出应答值，并存在负载中
			MessageDigest md5 = MessageDigest.getInstance("md5");
			md5.update(buf);
			payload = md5.digest();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	/**
	 * 计算key值
	 * 
	 * @param key
	 * @return
	 */
	private long calcKey(String key) {
		// 连接字符串中的数组
		StringBuffer numBuf = new StringBuffer();
		// 记录字符串中的个数
		int spaceCount = 0;
		for(int i = 0; i < key.length(); i++) {
			char c = key.charAt(i);
			if(c >= '0' && c <= '9')
				numBuf.append(c);
			else if(c == ' ')
				spaceCount ++;
		}

		// 上述二者相除，即为key值
		long num = Long.parseLong(numBuf.toString());
		if(spaceCount > 0)
			num /= spaceCount;

		return num;
	}

	/**
	 * 为正式版握手请求创建应答
	 * 
	 * @param request
	 */
	private void makeResponseForVer13(HandShakeRequest request) {
		// 标准头部
		headers.put(HEADER_UPGRADE, "websocket");
		headers.put(HEADER_CONNECTION, "Upgrade");
		headers.put(HEADER_WEBSOCKET_ORIGIN, request.getOrigin());
		headers.put(HEADER_WEBSOCKET_LOCATION, request.getLocation());

		// 计算返回密钥
		String key = request.getHeader("Sec-WebSocket-Key");
		String accept = key + SERVER_UUID;
		try {
			MessageDigest sha1 = MessageDigest.getInstance("sha1");
			byte[] buf = sha1.digest(accept.getBytes());
			accept = Base64.encode(buf);
		} catch (Exception e) {
			e.printStackTrace();
		}
		headers.put(HEADER_WEBSOCKET_ACCEPT, accept);
	}

	/**
	 * 将握手响应写到输出流中
	 * 
	 * @param out
	 */
	public void write(OutputStream out) {
		try {
			StringBuffer buf = new StringBuffer();
			buf.append(RESPONSE_STATE).append(CRLF);
			for(String key : headers.keySet()) {
				String value = headers.get(key);
				buf.append(key).append(": ").append(value).append(CRLF);
			}
			buf.append(CRLF);
			out.write(buf.toString().getBytes());

			if(payload != null) out.write(payload);
			out.flush();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
}
